import io
import re
import sys
import traceback
import zipfile

import numpy as np
import os
from pathlib import Path
import random
import gc

# 从 scikit-learn 中导入 train_test_split 函数，用于划分训练集和测试集。
from sklearn.model_selection import train_test_split


def preprocess_bed_to_npy(plant_name, dataset_id, files, flag):
    """ 通用数据处理函数
        Args:
            plant_name: 植物种类名称
            dataset_id: 数据集ID
            files: 上传成功但未处理的文件
            flag: 0 = 训练处理
                  1 = 预测处理
        Returns:
            result: 处理结果: True = 成功; False = 失败
            message: 处理信息
            chrom_count: 染色体种类数量
            chrom_str: 染色体名称组成的字符串
            file_size_mb: 文件大小（MB）
            sample_count: 样本数量
        """
    chrom_count = 0  # 染色体种类数量
    chrom_str = ""  # 染色体名称组成的字符串
    file_size_mb = ""  # 文件大小（MB）
    sample_count = 0  # 样本数量
    try:
        # 获取项目根目录
        current_file_path = os.path.abspath(__file__)  # 预期：F:\Project\PlantACR\ACRmodel\preprocess.py
        model_root = os.path.dirname(os.path.dirname(current_file_path))
        print(f"项目根目录：{model_root}")  # 预期：F:\Project\PlantACR

        # 处理输入文件列表
        input_files = [Path(model_root) / "ACRmodel" / "static" / "temp" / f.strip()
                       for f in files.split(',')]
        print("文件列表：", input_files)

        genome_file = Path(model_root) / "ACRmodel" / "static" / "fasta" / f"{plant_name}.fa"
        # 基于dataset_id创建目录
        # 确保唯一性，方便后续查找和管理
        # 避免用户修改数据集名字后对应目录名字修改出错，导致找不到数据集
        if flag == 0:
            # 训练处理
            save_dir = Path(model_root) / "ACRmodel" / "static" / "dataset" / plant_name / str(dataset_id)
        else:
            # 预测处理
            save_dir = Path(model_root) / "ACRmodel" / "static" / "predict_data" / plant_name / str(dataset_id)
        # 保证npy与bed的保存目录存在
        os.makedirs(save_dir, exist_ok=True)

        print("\033[1;30;34m%s\033[0m" % "合并数据集文件中...")
        sys.stdout.flush()  # 立即输出缓冲区
        output_file = Path(save_dir) / f"{dataset_id}.bed"
        result, chrom_count, chrom_str = merge_and_sort_bed_files(input_files, output_file)
        if not result:
            print("\033[1;30;34m%s\033[0m" % "合并数据集文件失败")
            result = False
            message = "文件合并失败"
            return result, message, chrom_count, chrom_str, file_size_mb, sample_count
        print("\033[1;30;34m%s\033[0m" % "合并数据集文件完成")

        print("\033[1;30;34m%s\033[0m" % "生成数据集样本文件中...")
        sys.stdout.flush()  # 立即输出缓冲区
        # 读取合并后的bed文件
        bed_file = output_file
        nuc = open(bed_file, 'r')
        bed = nuc.readlines()  # 读取bed文件，存储到变量 bed 中

        # 构建全基因字典 genome_fasta
        genome_fasta = {}  # 创建dict字典数据类型，用于存储需要的基因组序列数据
        seq_name = None  # 初始化seq_name
        # 打开并逐行读取基因组文件（fasta格式）
        for line in open(genome_file, 'r'):
            # 检查当前行是否为描述基因序列的标题行。在 FASTA 文件中，通常用 '>' 来标识基因序列的标题行。
            # >染色体名 dna:scaffold scaffold:IRGSP-1.0:AP008247:1:157458:1 REF
            # AAAGCAGAATACCGACTTGGCCCATGCGAGCCCATTTTT......
            # ......
            if line[0] == '>':
                # 带有 > 的行，从 > 到第一个空格之间的内容赋值为name，不包含 > 和 空格
                seq_name = line.lstrip('>').split(' ')[0]
                # 将基因的名称（seq_name）作为字典的键
                genome_fasta[seq_name] = []
            else:
                # 将当前行（已经移除了行首和行尾的空格）添加到对应名字的基因序列字典中。
                genome_fasta[seq_name].append(line.strip())
        # Python字典的值可以是任意类型，这里直接实现从列表到字符串的转换
        for keys, val in genome_fasta.items():  # 返回遍历的值
            genome_fasta[keys] = ''.join(val)  # 以空字符串分隔，把列表变为一行字符串

        '''
            创建仅存储了需要的基因组序列的 正样本fa文件： output_pos.fa
            格式: >1:26943-27251\n对应碱基序列
        '''
        filtered_lines = []  # 存储可以读取的序列为新的bed文件
        skip_chr_count = {}  # 创建跳过字典，存储跳过的染色体号及对应记录数
        col_num = 0  # 初始化一个计数器
        output_pos_fa_file = Path(save_dir) / "output_pos.fa"
        output_pos_fa = open(output_pos_fa_file, 'w')
        sample_count = len(bed)  # 计算染色质可访问区域文件的行数
        print("样本数: ", sample_count)
        # 利用输入的bed文件，获取每行对应的DNA序列，存入新的fa文件中
        for line in bed:
            col_num += 1  # 行数
            chr_name = line.rstrip().split()[0]  # 从当前行中获取染色体号
            # 检查染色体号是否在字典中存在对应的键
            if chr_name not in genome_fasta:  # 如果染色体号在字典中不存在对应的键
                if chr_name not in skip_chr_count:
                    skip_chr_count[chr_name] = 1  # 如果染色体号不在跳过字典中，则初始化为1
                else:
                    skip_chr_count[chr_name] += 1  # 如果染色体号在跳过字典中，则记录数加1
                continue
            else:  # 如果染色体号在字典中存在对应的键,则保留
                filtered_lines.append(line)
            # "peak"指的是基因组中的一个特定区域，通常指的是一个DNA序列
            start_end = line.rstrip().split()[1:]  # 从当前行中获取该 peak 的起始和终止位置
            # 将基因序列写入输出文件。它首先写入 peak 的标题行（'>' 开头），然后写入基因序列
            output_pos_fa.write('>' + chr_name + ':' + start_end[0] + '-' + start_end[1] + '\n' +
                                genome_fasta[chr_name][int(start_end[0]):int(start_end[1])] + '\n')
        output_pos_fa.close()
        pass_num = 0  # bed中跳过的数量
        # 输出全基因字典genome_fasta中不存在的染色体和对应的peak数
        if skip_chr_count:
            # print("已有基因组库中没有以下染色体，相关信息如下: \n")
            for chr_name, count in skip_chr_count.items():
                # print("染色体{}共{}个peak\n".format(chr_name, count))
                pass_num += count
            print("部分染色质在已有基因组库中无记录，总共跳过peak=" + str(pass_num) + "个")

        '''
            创建了output_pos.fa基因组序列的 正样本txt文件： output_pos.txt
            这些序列被写入一个新的 txt 文本文件中，每一行都以 '1' 开头，表示正样本，后面跟着序列本身
            txt中不存储染色体号，只存储序列本身
        '''
        bed_dict_pos = {}
        output_pos_txt_file = Path(save_dir) / "output_pos.txt"
        output_pos_txt = open(output_pos_txt_file, 'w')
        for line in open(output_pos_fa_file, 'r'):
            # 格式: >1:26943-27251\n对应碱基序列
            if line[0] == ">":
                seq_name = line.strip()
                bed_dict_pos[seq_name] = []
            else:
                bed_dict_pos[seq_name].append(line.strip())
        for keys, val in bed_dict_pos.items():
            bed_dict_pos[keys] = ''.join(val)
            output_pos_txt.write(f"1\t{bed_dict_pos[keys]}\n")  # 写入格式化行
        output_pos_txt.close()

        '''
            训练相关数据集处理需要负样本
        '''
        output_neg_fa_file = Path(save_dir) / "output_neg.fa"
        output_neg_txt_file = Path(save_dir) / "output_neg.txt"
        if flag == 0:
            '''
                创建基因组序列的正样本output_pos.fa对应的负样本fa文件： output_neg.fa
                用已知正样本序列随机打乱生成负样本序列比自行生成更方便
                保证正样本和负样本碱基序列的数量、长度一致，方便训练
                格式: >1\n对应碱基序列
            '''
            neg_fasta = {}
            output_neg_fa = open(output_neg_fa_file, 'w')
            for line in open(output_pos_fa_file, 'r'):
                if line[0] == '>':
                    seq_name = line.lstrip('>').strip()
                    neg_fasta[seq_name] = []
                else:
                    neg_fasta[seq_name].append(line.strip())
            for keys, val in neg_fasta.items():
                neg_fasta[keys] = ''.join(val)
                fa = list(neg_fasta[keys])
                random.shuffle(fa)  # 将碱基打乱顺序
                output_neg_fa.write(">" + keys + "\n" + "".join(fa) + "\n")
            output_neg_fa.close()

            '''
                创建了output_neg.fa基因组序列的 负样本txt文件： output_neg.txt
                这些序列被写入一个新的 txt 文本文件中，每一行都以 '0' 开头，表示负样本，后面跟着序列本身
                txt中不存储染色体号，只存储序列本身
            '''
            bed_dict_neg = {}
            neg_o = open(output_neg_txt_file, 'w')
            for line in open(output_neg_fa_file, 'r'):
                if line[0] == ">":
                    seq_name = line.strip()
                    bed_dict_neg[seq_name] = []
                else:
                    bed_dict_neg[seq_name].append(line.strip())
            for keys, val in bed_dict_neg.items():
                bed_dict_neg[keys] = ''.join(val)
                neg_o.write(f"0\t{bed_dict_neg[keys]}\n")
            neg_o.close()
            print("\033[1;30;34m%s\033[0m" % "生成数据集样本文件完成")

        '''
            使用“一次热编码”将每个序列预处理为4列的二维矩阵
            输入样本数据集txt文件：output_pos.txt、output_neg.txt
            序列中的每个碱基都被矢量化
            A|a -> [1,0,0,0]
            C|c -> [0,1,0,0]
            G|g -> [0,0,1,0]
            T|t -> [0,0,0,1]
            N|n -> [0,0,0,0]
        '''
        def load_data(sample_file):
            f = open(sample_file, 'r')
            sequences = f.readlines()
            num = len(sequences)  # 样本数
            data = np.empty((num, 1000, 4), dtype='float32')  # 三维数组：样本×碱基位置×特征
            label = np.empty((num,), dtype="int")  # 样本标签
            for i in range(num):
                lline = sequences[i].replace('\n', '')  # 移除行尾换行符
                list_line = re.split(r'\s+', lline)  # 按一个或多个空白字符（包括空格/tab）分割
                label[i] = list_line[0]  # 分割结果中的标签部分
                one_sequence = list_line[1]  # 分割结果中的DNA序列部分
                for j in range(1000):  # 只取前1000个碱基，不足的用N补
                    if j <= len(one_sequence) - 1:
                        # 编码，将类别型数据（碱基类型）转换为数值型数据
                        if re.findall(one_sequence[j], 'A|a'):
                            data[i, j, :] = np.array([1.0, 0.0, 0.0, 0.0], dtype='float32')
                        if re.findall(one_sequence[j], 'C|c'):
                            data[i, j, :] = np.array([0.0, 1.0, 0.0, 0.0], dtype='float32')
                        if re.findall(one_sequence[j], 'G|g'):
                            data[i, j, :] = np.array([0.0, 0.0, 1.0, 0.0], dtype='float32')
                        if re.findall(one_sequence[j], 'T|t'):
                            data[i, j, :] = np.array([0.0, 0.0, 0.0, 1.0], dtype='float32')
                        if re.findall(one_sequence[j], 'N|n'):
                            data[i, j, :] = np.array([0.0, 0.0, 0.0, 0.0], dtype='float32')
                    else:
                        data[i, j, :] = np.array([0.0, 0.0, 0.0, 0.0], dtype='float32')
            return data, label

        print("\033[1;30;34m%s\033[0m" % "生成数据集编码文件中...")
        # 调用一次热编码函数
        if flag == 0:
            # 训练处理
            sys.stdout.flush()  # 立即输出缓冲区
            data_pos, label_pos = load_data(output_pos_txt_file)  # 正数据集 + 对应标签
            data_neg, label_neg = load_data(output_neg_txt_file)  # 负数据集 + 对应标签

            # 根据60%、20%和20%将正数据集划分为训练集、验证集和测试集
            data_pos_train, data_pos_test, label_pos_train, label_pos_test = train_test_split(data_pos, label_pos,
                                                                                              test_size=0.4,
                                                                                              random_state=1)
            data_pos_test, data_pos_val, label_pos_test, label_pos_val = train_test_split(data_pos_test, label_pos_test,
                                                                                          test_size=0.5, random_state=1)
            # 根据60%、20%和20%将负数据集划分为训练集、验证集和测试集
            data_neg_train, data_neg_test, label_neg_train, label_neg_test = train_test_split(data_neg, label_neg,
                                                                                              test_size=0.4,
                                                                                              random_state=1)
            data_neg_test, data_neg_val, label_neg_test, label_neg_val = train_test_split(data_neg_test, label_neg_test,
                                                                                          test_size=0.5, random_state=1)

            # np.concatenate 函数来沿着指定的轴将两个或多个数组连接在一起。这里，axis=0 意味着沿着第一个轴（样本为个体）进行拼接。
            # 对于训练集，将正样本的训练集 data_pos_train 和负样本的训练集 data_neg_train 拼接在一起，得到合并后的训练集 data_train。同样地，将对应的标签集也进行拼接，得到 label_train。
            # 对于验证集，将正样本的验证集 data_pos_val 和负样本的验证集 data_neg_val 拼接在一起，得到合并后的验证集 data_val。同样地，将对应的标签集也进行拼接，得到 label_val。
            # 对于测试集，将正样本的测试集 data_pos_test 和负样本的测试集 data_neg_test 拼接在一起，得到合并后的测试集 data_test。同样地，将对应的标签集也进行拼接，得到 label_test。
            # 训练集
            data_train = np.concatenate((data_pos_train, data_neg_train), axis=0)  # 多个数组的拼接
            label_train = np.concatenate((label_pos_train, label_neg_train), axis=0)
            # 验证集
            data_val = np.concatenate((data_pos_val, data_neg_val), axis=0)
            label_val = np.concatenate((label_pos_val, label_neg_val), axis=0)
            # 测试集
            data_test = np.concatenate((data_pos_test, data_neg_test), axis=0)
            label_test = np.concatenate((label_pos_test, label_neg_test), axis=0)
            # 保存数据到指定路径
            np.save(Path(save_dir) / 'data_train.npy', data_train)
            np.save(Path(save_dir) / 'label_train.npy', label_train)
            np.save(Path(save_dir) / 'data_val.npy', data_val)
            np.save(Path(save_dir) / 'label_val.npy', label_val)
            np.save(Path(save_dir) / 'data_test.npy', data_test)
            np.save(Path(save_dir) / 'label_test.npy', label_test)
            print("训练数据处理完成")
            result = True
            message = "数据处理完成"
            print("\033[1;30;34m%s\033[0m" % "生成数据集编码文件完成")
            # 计算目录下所有文件总大小
            total_size = sum(os.path.getsize(os.path.join(save_dir, f)) for f in os.listdir(save_dir)
                             if os.path.isfile(os.path.join(save_dir, f)))
            file_size_mb = f"{total_size / (1024 * 1024):.2f}MB"
            print(f"数据集总大小: {file_size_mb}")
            # 将所有处理好的文件打包成ZIP
            zip_buffer = io.BytesIO()
            with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zipf:
                for file in os.listdir(save_dir):
                    file_path = os.path.join(save_dir, file)
                    if os.path.isfile(file_path):  # 只打包文件，不打包子目录
                        zipf.write(file_path, arcname=file)  # 保持文件名不变
            # 保存ZIP文件
            zip_filename = f"dataset_{dataset_id}.zip"
            zip_path = os.path.join(save_dir, zip_filename)
            with open(zip_path, 'wb') as f:
                f.write(zip_buffer.getvalue())
            print(f"已生成压缩包: {zip_path}")
        else:
            # 预测处理，仅处理正样本即可
            data_pos, label_pos = load_data(output_pos_txt_file)  # 正数据集 + 对应标签（不需要标签）
            # 保存数据到指定路径
            np.save(Path(save_dir) / 'data_predict.npy', data_pos)
            print("预测数据处理完成")
            result = True
            message = "预测数据处理完成"
            print("\033[1;30;34m%s\033[0m" % "生成数据集编码文件完成")

        gc.collect()  # 手动触发垃圾回收
        return result, message, chrom_count, chrom_str, file_size_mb, sample_count

    except Exception as e:
        print("预处理出错: ", e)
        # 使用 traceback 打印完整的错误信息和堆栈跟踪
        traceback.print_exc()
        result = False
        message = "数据处理失败"
        gc.collect()  # 手动触发垃圾回收
        return result, message, chrom_count, chrom_str, file_size_mb, sample_count


def merge_and_sort_bed_files(input_files, output_file):
    """
    Args:
        input_files: 需要合并的文件列表
        output_file: 合并后的输出文件

    Returns:
        合并成功返回True，否则返回False
        chrom_count: 染色体种类数量
        chrom_str: 染色体名称组成的字符串
    """
    # 统计染色体信息
    chrom_set = set()  # 用于存储不重复的染色体名称
    chrom_names = []  # 用于按顺序存储染色体名称
    chrom_count = 0  # 染色体种类数量
    chrom_str = ""  # 染色体名称组成的字符串
    try:
        init_intervals = []
        # 合并所有区间
        for file_path in input_files:
            with open(file_path, 'r') as file:
                for line in file:
                    line = line.strip()
                    if not line:  # 跳过空行
                        print("跳过空")
                        continue
                    # 拆分所有字段，取前三个
                    parts = line.split('\t')
                    if len(parts) < 3:  # 至少需要3列数据
                        print("小于3")
                        continue
                    chrom, start_str, end_str = parts[:3]
                    # 检查start和end是否为整数
                    try:
                        start = int(start_str)
                        end = int(end_str)
                    except ValueError:
                        # 非整数行跳过（如标题行）
                        print("start_str: ", start_str, "end_str: ", end_str)
                        print("格式错误")
                        continue
                    init_intervals.append((chrom, start, end))
        # 按照染色体名称和起始位置排序
        sorted_intervals = sorted(init_intervals, key=lambda x: (x[0], x[1]))

        # 区间合并
        merged_intervals = []  # 最终合并结果
        current_chrom = ''  # 当前处理的染色体名称（初始为空）
        current_start = float('inf')  # 当前区间起始位置（初始为正无穷，确保第一个区间会被处理）
        current_end = float('-inf')  # 当前区间结束位置（初始为负无穷）

        for chrom, start, end in sorted_intervals:
            # 统计染色体信息
            if chrom not in chrom_set:
                chrom_set.add(chrom)
                chrom_names.append(chrom)
            # 染色体不同 或 新区间的起始位置 > 当前区间的结束位置（无重叠/不相邻）
            if chrom != current_chrom or start > current_end:
                if current_chrom:  # 非初始状态，说明有已处理的区间
                    merged_intervals.append((current_chrom, current_start, current_end))
                # 重置当前区间为当前遍历的区间
                current_chrom = chrom
                current_start = start
                current_end = end
            else:
                # 染色体相同 且 新区间的起始位置 ≤ 当前区间的结束位置（有重叠/相邻）
                current_end = max(current_end, end)
        # 处理最后一个区间
        if current_chrom:
            merged_intervals.append((current_chrom, current_start, current_end))

        # 计算染色体统计信息
        chrom_count = len(chrom_set)  # 染色体种类数量
        chrom_str = ",".join(chrom_names)  # 染色体名称组成的字符串

        # 将合并结果写入输出文件
        with open(output_file, 'w') as file:
            for chrom, start, end in merged_intervals:
                file.write(f"{chrom}\t{start}\t{end}\n")
        print("合并完成")
        return True, chrom_count, chrom_str
    except Exception as e:
        print("合并错误: ", e)
        return False, chrom_count, chrom_str
